Skip to content

【训练营】AI Compiler 项目#20

Open
ArcaLunar wants to merge 16 commits intobitzyz:mainfrom
ArcaLunar:main
Open

【训练营】AI Compiler 项目#20
ArcaLunar wants to merge 16 commits intobitzyz:mainfrom
ArcaLunar:main

Conversation

@ArcaLunar
Copy link
Copy Markdown

算子添加

  1. 为框架添加基于 InfiniCore 的 kerneloperator 算子适配,包括 clip, conv, rms_norm 等等
  2. unified_converter.py 中为算子进行 PyTorch 注册
  3. 为算子添加 kernel tests 和 operator tests

Test 通过截图

image

图优化功能支持

  1. 图优化显式触发,支持原地改写;通过继承 GraphPass 类可以实现其他的图优化策略并统一管理
  2. 添加默认优化流水线 (graph_optimizer.cc):拓扑、形状推理、常量折叠、Identity 消除、幂等消除、死代码消除等等
  3. 为图优化添加必要的测试:test_graph_optimizer

CUDA Graph 支持

  1. 通过 setCudaGraphEnabled() 显式启用,且要求 Graph 在 CUDA 设备上
  2. 通过 infinirt 框架支持对 Graph 进行捕获并将 Graph 编译为 CUDA Graph(结果保存为线程本地缓存),以及 CUDA Graph 的运行
  3. 默认在 Graph 改变时调用 invalidateCudaGraph() 使缓存的 CUDA Graph 失效
  4. test_graph.cc 中添加简单测试:CudaGraphLifecycleBookkeeping, CudaGraphInvalidatesOnOutputMutation, CudaGraphCompileLaunchAndRecompileAfterMutation

限制

  1. 多设备场景
  2. 未实现增量更新

@ArcaLunar ArcaLunar marked this pull request as ready for review March 17, 2026 02:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants